import requests
import json
import base64
import os
from PIL import Image
import io
from datetime import datetime

# 中转API配置
BASE_URL = "https://www.chataiapi.com/v1"
API_KEY = "sk-K4KYV1kGQ3C4DN9ZZMlG84oJY1LN6zG0xbMGm1pidiqav1Sc"

def encode_image_to_base64(image_path):
    """将图片文件编码为base64字符串"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def decode_base64_to_image(base64_string, output_path):
    """将base64字符串解码并保存为图片文件"""
    image_data = base64.b64decode(base64_string)
    with open(output_path, "wb") as image_file:
        image_file.write(image_data)

def generate_image_with_text(prompt, model_name="gemini-2.0-flash-exp"):
    """使用文本提示生成图片"""
    try:
        url = f"{BASE_URL}/chat/completions"
        
        payload = json.dumps({
            "model": model_name,
            "messages": [
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "response_format": {"type": "image_url"}
        })
        
        headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {API_KEY}',
            'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
            'Content-Type': 'application/json'
        }
        
        response = requests.post(url, headers=headers, data=payload)
        
        if response.status_code == 200:
            result = response.json()
            print("API调用成功！")
            
            # 检查响应结构
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    content = choice['message']['content']
                    
                    # 检查是否是图片内容
                    if isinstance(content, str) and content.startswith('data:image/'):
                        # 提取base64数据
                        if ',' in content:
                            base64_data = content.split(',')[1]
                            return base64.b64decode(base64_data)
                    else:
                        print("响应内容不是图片格式")
                        print("响应内容:", content)
            else:
                print("未找到有效的响应内容")
                print("完整响应:", json.dumps(result, indent=2, ensure_ascii=False))
        else:
            print(f"API调用失败，状态码: {response.status_code}")
            print("错误信息:", response.text)
            
        return None
            
    except Exception as e:
        print(f"生成图片时出错: {e}")
        return None

def understand_image(image_path, prompt="请描述这张图片"):
    """理解图片内容"""
    try:
        url = f"{BASE_URL}/chat/completions"
        
        # 编码图片为base64
        image_base64 = encode_image_to_base64(image_path)
        
        payload = json.dumps({
            "model": "gemini-2.0-flash-exp",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image_base64}"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 1000
        })
        
        headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {API_KEY}',
            'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
            'Content-Type': 'application/json'
        }
        
        response = requests.post(url, headers=headers, data=payload)
        
        if response.status_code == 200:
            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    return choice['message']['content']
        else:
            print(f"API调用失败，状态码: {response.status_code}")
            print("错误信息:", response.text)
            
        return None
        
    except Exception as e:
        print(f"理解图片时出错: {e}")
        return None

def generate_image_from_images(image1_path, image2_path, prompt="请基于这两张图片生成一张新的图片", model_name="gemini-2.0-flash-exp"):
    """基于两张图片生成新图片"""
    try:
        url = f"{BASE_URL}/chat/completions"
        
        # 编码两张图片为base64
        image1_base64 = encode_image_to_base64(image1_path)
        image2_base64 = encode_image_to_base64(image2_path)
        
        payload = json.dumps({
            "model": model_name,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image1_base64}"
                            }
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image2_base64}"
                            }
                        }
                    ]
                }
            ],
            "response_format": {"type": "image_url"}
        })
        
        headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {API_KEY}',
            'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
            'Content-Type': 'application/json'
        }
        
        response = requests.post(url, headers=headers, data=payload)
        
        if response.status_code == 200:
            result = response.json()
            print("API调用成功！")
            
            # 检查响应结构
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    content = choice['message']['content']
                    
                    # 检查是否是图片内容
                    if isinstance(content, str) and content.startswith('data:image/'):
                        # 提取base64数据
                        if ',' in content:
                            base64_data = content.split(',')[1]
                            return base64.b64decode(base64_data)
                    else:
                        print("响应内容不是图片格式")
                        print("响应内容:", content)
            else:
                print("未找到有效的响应内容")
                print("完整响应:", json.dumps(result, indent=2, ensure_ascii=False))
        else:
            print(f"API调用失败，状态码: {response.status_code}")
            print("错误信息:", response.text)
            
        return None
            
    except Exception as e:
        print(f"生成图片时出错: {e}")
        return None

def save_image_data(image_data, output_path):
    """保存图片数据到文件"""
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 将图片数据写入文件
        with open(output_path, "wb") as f:
            f.write(image_data)
        
        print(f"图片已保存到: {output_path}")
        return output_path
        
    except Exception as e:
        print(f"保存图片时出错: {e}")
        return None

def process_images_with_boots(girl_image_path, boots_image_path, prompt="图中女生穿上牛仔裤和图二中靴子, 保持脸部不变, 镜头拉远保证女生全身和靴子均在图内", api_key=None):
    """
    处理女生图片和靴子图片，生成新的合成图片
    
    Args:
        girl_image_path (str): 女生图片路径
        boots_image_path (str): 靴子图片路径
        prompt (str): 提示词
        api_key (str): API密钥，如果不提供则使用默认密钥
    
    Returns:
        str: 生成的图片路径，如果失败返回None
    """
    # 检查输入文件是否存在
    if not os.path.exists(girl_image_path):
        print(f"错误: 女生图片文件不存在: {girl_image_path}")
        return None
    
    if not os.path.exists(boots_image_path):
        print(f"错误: 靴子图片文件不存在: {boots_image_path}")
        return None
    
    # 生成图片
    print("正在调用中转API生成图片...")
    image_data = generate_image_from_images(girl_image_path, boots_image_path, prompt)
    
    if image_data:
        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(girl_image_path))[0]
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"{base_name}_ai_{timestamp}.png"
        output_path = os.path.join("pic", output_filename)
        
        # 保存图片
        result_path = save_image_data(image_data, output_path)
        if result_path:
            print(f"图片生成完成！输出文件: {result_path}")
            return result_path
        else:
            print("图片保存失败")
            return None
    else:
        print("图片生成失败")
        return None

def generate_text_to_image(prompt, output_dir="pic", api_key=None):
    """根据文本提示生成图片"""
    print(f"正在根据提示生成图片: {prompt}")
    image_data = generate_image_with_text(prompt)
    
    if image_data:
        # 生成输出文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"generated_{timestamp}.png"
        output_path = os.path.join(output_dir, output_filename)
        
        # 保存图片
        result_path = save_image_data(image_data, output_path)
        if result_path:
            print(f"图片生成完成！输出文件: {result_path}")
            return result_path
        else:
            print("图片保存失败")
            return None
    else:
        print("图片生成失败")
        return None

def analyze_image(image_path, analysis_prompt="请详细分析这张图片的内容", api_key=None):
    """分析图片内容"""
    print(f"正在分析图片: {image_path}")
    result = understand_image(image_path, analysis_prompt)
    
    if result:
        print("图片分析结果:")
        print(result)
        return result
    else:
        print("图片分析失败")
        return None

if __name__ == "__main__":
    # 示例使用
    # 注意：使用中转API，无需魔法上网
    
    # 示例1: 基于两张图片生成新图片
    girl_image = "pic/2024-08-06_19.52.34_比比好可爱_一日MISS_DIOR体验卡_6.png"
    boots_image = "src/boots.png"
    
    # 使用默认提示词
    result = process_images_with_boots(girl_image, boots_image)
    
    # 示例2: 根据文本生成图片
    # text_prompt = "一只可爱的小猫坐在花园里"
    # result = generate_text_to_image(text_prompt)
    
    # 示例3: 分析图片
    # analysis_result = analyze_image(girl_image, "请描述这张图片中的人物和场景")
